import asyncio
import json
import os
import random

from typing import List

from agents import Agent, RunContextWrapper, Runner, function_tool, handoff
from agents.models.openai_provider import OpenAIProvider
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX

from game_helpers import get_random_card_list, get_cached_expressions


API_KEY = os.getenv("EXAMPLE_API_KEY")
BASE_URL = os.getenv("EXAMPLE_BASE_URL")

provider = OpenAIProvider(
        api_key=API_KEY,
        base_url=BASE_URL,
        use_responses=False
    )
model = provider.get_model("qwen2.5")


instruction_template_dict = {
    "GameDealer": """
        Generate 4 random natural numbers between 1 and 13, include 1 and 13. Just return 4 numbers in an array, don't include other content. The returned array should not be repeated with the following arrays:
        {old_arrays}
        """,
    "MathProdigy": """You are a helpful agent. """,
    "GameJudger": """You are a helpful agent. """,
    "GamePlayer": """You are a helpful agent. """,
}


user_prompt_template_dict = {
    "GameDealer": """generate an array""",
    "MathProdigy": """
        What's the 24 points expression of {last_cards_posted} ? 
        If the result is 'expression not found', just return 'expression not found'. 
        If the result is an arithmetic expression, return the expression follow this Python dictionary format:{{"expression": "result"}}, just return this Python dictionary and do not add anything else.
        """,
    "GameJudger": """
        Cards posted is '{last_cards_posted}', what's the check result of {expression} ? Return the check result follow this Python dictionary format:{{"check_result": "result"}}, the result shoule be 'Correct' or 'Wrong', just return this Python dictionary and do not add anything else. 
    """,
    "GamePlayer": """
        What's the human user's reply of {last_cards_posted} ? Return the reply follow this Python dictionary format:{{"human_reply": "reply"}}, just return this Python dictionary and do not add anything else.
    """,
}


class CustomContext:
    def __init__(self, context_variables: dict):
        self.context_variables = context_variables


def get_instruction(
    run_context: RunContextWrapper[CustomContext], agent: Agent[CustomContext]
) -> str:
    global instruction_template_dict

    context = run_context.context
    context_variables = context.context_variables

    agent_name = context_variables["agent_name"]
    instruction_template = instruction_template_dict[agent_name]
    instruction = instruction_template

    if agent_name == "GameDealer":
        last_cards_posted = context_variables["old_arrays"]
        instruction = instruction_template.format(old_arrays=last_cards_posted)
   
    return instruction


def get_user_prompt(context_variables):
    global user_prompt_template_dict

    agent_name = context_variables["agent_name"]
    user_prompt_template = user_prompt_template_dict[agent_name]
    user_prompt = user_prompt_template

    if agent_name == "MathProdigy":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    elif agent_name == "GameJudger":
        expression = context_variables["expression"]
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(expression=expression,last_cards_posted=last_cards_posted)

    elif agent_name == "GamePlayer":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    return user_prompt


@function_tool
def get_24_points_expression_func(last_cards_posted: str) -> str:
    """Resolve the expression of 24 points game, return an arithmetic expression.

    Keyword arguments:
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    point_list = json.loads(last_cards_posted)  
    if len(point_list) == 0:
        return "expression not found"

    expressions = get_cached_expressions(point_list)
    
    result = "expression not found"
    if len(expressions) > 0:
        random_idx = random.randint(0, len(expressions)-1)
        expression = f"'{expressions[random_idx]}'".replace("'", "")
        print(f"The resolved 24 points expression is '{expression}'")
        return expression
    
    return "expression not found"


@function_tool
def check_24_points_expression_func(expression: str, last_cards_posted: str) -> str:
    """Check if the result of an arithmetic expression is equal 24, return 'Correct' or 'Wrong'.

    Keyword arguments:
      expression: an arithmetic expression
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    result = eval(expression.replace("'", ""))
    if abs(result - 24) < 0.001:
        return "Correct"

    return "Wrong"
    

@function_tool
def get_human_reply_func(last_cards_posted: str) -> str:
    """Interact with a human user and get a reply from the human user. The replay should be 'deal', 'help', 'exit' or an an arithmetic expression.

    Keyword arguments:
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    PROMPT_TEMPLATE: str = """
    Cards the dealer just posted: {content}
    Please give an expression for the four operations that results in 24.
    Type 'help' if you feel it's difficult.
    Type 'deal' if you want the dealer to deal cards again.
    Type 'exit' if you want to exit this game, type 'exit'.
    """

    point_list = json.loads(last_cards_posted)
    card_list = get_random_card_list(point_list)
    cards_content = f"{{'card_list': {card_list}, 'point_list': {point_list}}}"

    prompt = PROMPT_TEMPLATE.format(content=cards_content)
    human_reply = input(prompt)
    
    return human_reply


agent_bill = Agent(
    name="GameDealer",
    instructions=get_instruction,
    model=model,
    tools=[],
)

agent_gauss = Agent(
    name="MathProdigy",
    handoff_description="A math prodigy who can help users",    
    instructions=get_instruction,
    model=model,
    tools=[get_24_points_expression_func],
)

agent_peter = Agent(
    name="GameJudger",
    handoff_description="Person responsible for checking if the expression is correct",     
    instructions=get_instruction,
    model=model,
    tools=[check_24_points_expression_func],
)

agent_david = Agent(
    name="GamePlayer",
    instructions=get_instruction,
    model=model,
    tools=[get_human_reply_func],
)

async def on_handoff_gauss(ctx: RunContextWrapper[None]):
    # print(f"in on_handoff_gauss(), context_variables is {ctx.context.context_variables} ")
    pass

async def on_handoff_peter(ctx: RunContextWrapper[None]):
    # print(f"in on_handoff_peter(), context_variables is {ctx.context.context_variables} ")
    pass

handoff_gauss = handoff(
    agent=agent_gauss,
    on_handoff=on_handoff_gauss,
    tool_name_override="help_agent",
    tool_description_override="an agent who can help user to calculate 24 points expression",
)

handoff_peter = handoff(
    agent=agent_peter,
    on_handoff=on_handoff_peter,
    tool_name_override="judge_agent",
    tool_description_override="an agent who can judge whether the 24 points expression is correct",
)

triage_agent = Agent[CustomContext](
    name="Triage Agent",
    handoff_description="A triage agent that can delegate a customer's request to the appropriate agent.",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX} "
        "You are a helpful triaging agent. You can use your tools to delegate questions to other appropriate agents."
    ),
    handoffs=[
        handoff_gauss,
        handoff_peter,
    ],
)


async def deal_cards(old_arrays: List[int]) -> str:
    global agent_bill

    print(f"used old_arrays is :{old_arrays}")
    context_var_dict = {
        "agent_name":"GameDealer", 
        "old_arrays": f"{old_arrays}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_bill,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    cards_posted = response.final_output
    print(cards_posted)
    return cards_posted


async def run_game_one_turn(last_cards_posted: str) -> tuple:
    global agent_david
    
    context_var_dict = {
        "agent_name":"GamePlayer", 
        "last_cards_posted": f"{last_cards_posted}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_david,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    human_reply = eval(response.final_output)['human_reply']

    if human_reply == "deal" or human_reply == "exit":
        return human_reply, None

    result_expression = "expression not found"
    if human_reply == "help":
        context_var_dict = {
            "agent_name":"MathProdigy", 
            "last_cards_posted": f"{last_cards_posted}"
        }
        context = CustomContext(context_variables=context_var_dict)
        response = await Runner.run(
            triage_agent,
            [{"role": "user", "content": get_user_prompt(context_var_dict)}],
            context=context
        )

        result = response.final_output
        if "expression not found" in result:
            return "expression not found", "Correct"
        else:
            result_expression = eval(result)['expression']

    else:
        result_expression = human_reply

    context_var_dict = {
        "agent_name":"GameJudger", 
        "expression": f"{result_expression}",
        "last_cards_posted": f"{last_cards_posted}"
    }

    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        triage_agent,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    result = eval(response.final_output)['check_result']
    return result_expression, result


async def main_func():
    old_arrays = []
    last_cards_posted = await deal_cards(old_arrays)

    while True:
        result, check_result = await run_game_one_turn(last_cards_posted)
        print(f"In this turn, cards posted is '{last_cards_posted}', expression is '{result}', check result is '{check_result}'.")
        
        if result == "deal" or result == "expression not found" or check_result == "Correct":
            old_arrays.append(json.loads(last_cards_posted))
            last_cards_posted = await deal_cards(old_arrays)
 
        elif result == "exit":
            print("Bye bye, have a good day!")
            break
            
        
if __name__ == "__main__":
    asyncio.run(main_func())
